#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据可视化引擎核心模块
"""

import json
import asyncio
import time
import uuid
from typing import Dict, List, Any, Optional, Tuple, Union
import pandas as pd

from data_engine.utils.logger import get_logger, EngineLogger
from data_engine.utils.llm_adapter import LLMClient
from data_engine.core.chart_identifier import ChartIdentifier
from data_engine.core.data_processor import DataProcessor
from data_engine.core.output_generator import OutputGenerator
from data_engine.models.data_models import Dataset, QueryContext, VisualizationRequest, VisualizationResponse
from data_engine.config import ConfigManager

logger = get_logger(__name__)


class DataEngine:
    """数据可视化引擎主类"""
    
    def __init__(self, llm_client: LLMClient, 
                 small_data_threshold: int = 50,
                 enable_grouping_filtering: bool = True,
                 debug_mode: bool = False):
        """
        初始化数据引擎
        
        Args:
            llm_client: 已配置好的LLM客户端
            small_data_threshold: 小数据集阈值，默认50行
            enable_grouping_filtering: 是否启用分组过滤，默认True
            debug_mode: 是否启用调试模式，默认False
        """
        self.llm_client = llm_client
        self.config_manager = ConfigManager()
        
        # 优化配置参数
        self.small_data_threshold = small_data_threshold
        self.enable_grouping_filtering = enable_grouping_filtering
        self.debug_mode = debug_mode
        
        # 初始化引擎专用日志器
        self.engine_logger = EngineLogger("DataEngine")
        
        # 初始化处理器
        self.chart_identifier = ChartIdentifier(llm_client)
        self.data_processor = DataProcessor(llm_client, safe_execution=True)
        self.output_generator = OutputGenerator(llm_client)
        
        # 存储已加载的数据集
        self.datasets = {}
        
        logger.info(f"数据可视化引擎初始化完成 - 小数据阈值: {small_data_threshold}, 分组过滤: {enable_grouping_filtering}, 调试模式: {debug_mode}")
    
    def _log_debug(self, message: str, session_id: str = None) -> None:
        """统一的调试日志方法"""
        if not self.debug_mode:
            return
        
        if session_id:
            self.engine_logger.set_session_id(session_id)
        self.engine_logger.debug(message)
    
    def _log_info(self, message: str, session_id: str = None) -> None:
        """统一的信息日志方法"""
        if session_id:
            self.engine_logger.set_session_id(session_id)
        self.engine_logger.info(message)
    
    def _log_warning(self, message: str, session_id: str = None) -> None:
        """统一的警告日志方法"""
        if session_id:
            self.engine_logger.set_session_id(session_id)
        self.engine_logger.warning(message)
    
    def _log_error(self, message: str, session_id: str = None) -> None:
        """统一的错误日志方法"""
        if session_id:
            self.engine_logger.set_session_id(session_id)
        self.engine_logger.error(message)
    
    def _get_reasoning_info(self, reasoning: str) -> Optional[str]:
        """根据调试模式返回推理信息"""
        return reasoning if self.debug_mode else None
    
    def set_debug_mode(self, debug_mode: bool) -> None:
        """设置调试模式"""
        self.debug_mode = debug_mode
        self._log_info(f"调试模式已{'启用' if debug_mode else '禁用'}")
    
    def is_debug_mode(self) -> bool:
        """检查是否为调试模式"""
        return self.debug_mode
    
    async def _process_single_row_data(self, 
                                      dataset: Dataset, 
                                      question: str, 
                                      session_id: str, 
                                      start_time: float) -> VisualizationResponse:
        """处理单行数据的快速table生成"""
        try:
            self._log_info("开始单行数据快速处理", session_id)
            
            # 生成字段描述（使用AI）
            field_descriptions = await self._generate_field_descriptions_for_table(dataset, question)
            
            # 构建table配置
            columns = []
            for col in dataset.df.columns:
                col_desc = field_descriptions.get(col, col)
                col_type = "number" if pd.api.types.is_numeric_dtype(dataset.df[col]) else "string"
                columns.append({
                    "key": col,
                    "title": col_desc,
                    "dataType": col_type
                })
            
            chart_config = {
                "type": "table",
                "columns": columns,
                "data": dataset.df.to_dict('records'),
                "pagination": False,
                "pageSize": 10,
                "total": 1
            }
            
            # 创建简化的分析结果
            analysis_result = {
                "chart_type": "table",
                "need_chart": False,
                "x_field": None,
                "y_field": None,
                "y_fields": [],
                "reasoning": self._get_reasoning_info("单行数据自动选择表格展示")
            }
            
            execution_time = time.time() - start_time
            self._log_info(f"单行数据处理完成，耗时: {execution_time:.3f}秒", session_id)
            
            return VisualizationResponse(
                success=True,
                chart_config=chart_config,
                data=dataset.df.to_dict('records'),
                chart_type="table",
                execution_time=execution_time,
                session_id=session_id,
                analysis_result=analysis_result,
                message="单行数据快速处理完成"
            )
            
        except Exception as e:
            execution_time = time.time() - start_time
            error_msg = f"单行数据处理失败: {str(e)}"
            self._log_error(error_msg, session_id)
            
            return VisualizationResponse(
                success=False,
                chart_config=None,
                data=[],
                chart_type="table",
                execution_time=execution_time,
                session_id=session_id,
                error=error_msg,
                message=error_msg
            )
     
    async def _process_small_dataset_direct(self, 
                                           dataset: Dataset, 
                                           question: str, 
                                           chart_type: Optional[str], 
                                           session_id: str, 
                                           start_time: float) -> VisualizationResponse:
        """处理小数据集的快速图表生成"""
        try:
            self._log_info("开始小数据集快速处理", session_id)
            
            # 直接生成图表配置（使用AI）
            chart_config = await self._generate_chart_config_direct(dataset, question, chart_type)
            
            # 创建简化的分析结果
            analysis_result = {
                "chart_type": chart_config.get("type", "table"),
                "need_chart": chart_config.get("type", "table") != "table",
                "x_field": None,
                "y_field": None,
                "y_fields": [],
                "reasoning": self._get_reasoning_info("小数据集快速处理，无需复杂分析")
            }
            
            execution_time = time.time() - start_time
            self._log_info(f"小数据集处理完成，耗时: {execution_time:.3f}秒", session_id)
            
            return VisualizationResponse(
                success=True,
                chart_config=chart_config,
                data=dataset.df.to_dict('records'),
                chart_type=chart_config.get("type", "table"),
                execution_time=execution_time,
                session_id=session_id,
                analysis_result=analysis_result,
                message="小数据集快速处理完成"
            )
            
        except Exception as e:
            execution_time = time.time() - start_time
            error_msg = f"小数据集处理失败: {str(e)}"
            self._log_error(error_msg, session_id)
            
            return VisualizationResponse(
                success=False,
                chart_config=None,
                data=[],
                chart_type=chart_type or "table",
                execution_time=execution_time,
                session_id=session_id,
                error=error_msg,
                message=error_msg
            )
      
    async def _generate_field_descriptions_for_table(self, 
                                                    dataset: Dataset, 
                                                    question: str) -> Dict[str, str]:
        """为表格生成字段描述"""
        try:
            # 构建提示词
            prompt = f"""
根据用户问题和数据字段，为表格生成合适的中文字段描述。

用户问题：{question}

数据字段信息：
{self._get_field_info_for_prompt(dataset)}

请返回JSON格式的字段描述映射，格式如下：
{{
    "field_name": "中文描述",
    "field_name2": "中文描述2"
}}

要求：
1. 描述要简洁明了
2. 符合用户问题的上下文
3. 使用中文
4. 只返回JSON，不要其他内容
"""
            
            messages = [
                {"role": "user", "content": prompt}
            ]
            
            response = await self.llm_client.acompletion(messages)
            result_text = response.choices[0].message.content.strip()
            
            # 解析JSON
            if result_text.startswith('```json'):
                result_text = result_text.replace('```json', '').replace('```', '').strip()
            elif result_text.startswith('```'):
                result_text = result_text.replace('```', '').strip()
            
            field_descriptions = json.loads(result_text)
            return field_descriptions
            
        except Exception as e:
            self._log_warning(f"生成字段描述失败: {str(e)}，使用默认描述")
            # 返回默认描述
            return {col: col for col in dataset.df.columns}
    
    async def _generate_chart_config_direct(self, 
                                          dataset: Dataset, 
                                          question: str, 
                                          chart_type: Optional[str]) -> Dict[str, Any]:
        """直接生成图表配置"""
        try:
            # 构建提示词
            chart_type_hint = f"，图表类型建议为：{chart_type}" if chart_type else ""
            
            prompt = f"""
根据用户问题和数据，直接生成完整的图表配置JSON。

用户问题：{question}{chart_type_hint}

数据信息：
- 数据行数：{len(dataset.df)}
- 字段信息：
{self._get_field_info_for_prompt(dataset)}

数据样例（前3行）：
{dataset.df.head(3).to_string()}

请直接返回完整的图表配置JSON，支持的图表类型：table, bar, line, pie, scatter。

对于table类型，格式如下：
{{
    "type": "table",
    "columns": [
        {{"key": "字段名", "title": "中文标题", "dataType": "string/number"}}
    ],
    "data": [...],
    "pagination": true/false,
    "pageSize": 10,
    "total": 数据总数
}}

对于其他图表类型，请参考ECharts格式。

只返回JSON配置，不要其他内容。
"""
            
            messages = [
                {"role": "user", "content": prompt}
            ]
            
            response = await self.llm_client.acompletion(messages)
            result_text = response.choices[0].message.content.strip()
            
            # 解析JSON
            if result_text.startswith('```json'):
                result_text = result_text.replace('```json', '').replace('```', '').strip()
            elif result_text.startswith('```'):
                result_text = result_text.replace('```', '').strip()
            
            chart_config = json.loads(result_text)
            
            # 确保data字段包含实际数据
            if chart_config.get("type") == "table":
                chart_config["data"] = dataset.df.to_dict('records')
                chart_config["total"] = len(dataset.df)
            
            return chart_config
            
        except Exception as e:
            self._log_warning(f"生成图表配置失败: {str(e)}，使用默认table配置")
            # 返回默认table配置
            columns = []
            for col in dataset.df.columns:
                col_type = "number" if pd.api.types.is_numeric_dtype(dataset.df[col]) else "string"
                columns.append({
                    "key": col,
                    "title": col,
                    "dataType": col_type
                })
            
            return {
                "type": "table",
                "columns": columns,
                "data": dataset.df.to_dict('records'),
                "pagination": len(dataset.df) > 10,
                "pageSize": 10,
                "total": len(dataset.df)
            }
    
    def _get_field_info_for_prompt(self, dataset: Dataset) -> str:
        """获取字段信息用于提示词"""
        field_info = []
        for col in dataset.df.columns:
            dtype = str(dataset.df[col].dtype)
            unique_count = dataset.df[col].nunique()
            field_info.append(f"- {col}: {dtype}, {unique_count} unique values")
        return "\n".join(field_info)
      
    def load_dataset(self,
                    df: pd.DataFrame,
                    dataset_name: str = "default",
                    dataset_desc: str = None,
                    table_desc: str = None,
                    field_desc: Dict[str, str] = None) -> None:
        """
        加载数据集到引擎中
        
        Args:
            df: 数据DataFrame
            dataset_name: 数据集名称
            dataset_desc: 数据集描述（可选）
            table_desc: 表格描述（可选）
            field_desc: 字段描述字典（可选）
        """
        # 处理描述信息 - 保持原始状态，不强制生成默认值
        # 这样在提示词生成时可以根据实际情况决定是否使用
        
        # 标准化空字符串为None
        if dataset_desc == "":
            dataset_desc = None
        if table_desc == "":
            table_desc = None
        if field_desc == {}:
            field_desc = None
            
        # 如果提供了部分字段描述，保持原样，不自动补充
        # 这样AI可以根据数据内容和已有描述进行推理
        
        # 创建Dataset对象
        dataset = Dataset(
            dataset_id=f"dataset_{dataset_name}_{int(time.time())}",
            df=df,
            dataset_description=dataset_desc,
            table_description=table_desc,
            field_descriptions=field_desc
        )
        
        # 存储数据集
        self.datasets[dataset_name] = dataset
        
        # 记录加载信息
        desc_info = []
        if dataset_desc:
            desc_info.append("数据集描述")
        if table_desc:
            desc_info.append("表格描述")
        if field_desc:
            desc_info.append(f"字段描述({len(field_desc)}个)")
        
        desc_status = f"包含: {', '.join(desc_info)}" if desc_info else "无描述信息，将基于数据内容推理"
        self._log_info(f"数据集 '{dataset_name}' 加载完成: {len(df)} 行, {len(df.columns)} 列, {desc_status}")
    
    def _generate_default_field_descriptions(self, df: pd.DataFrame) -> Dict[str, str]:
        """生成默认字段描述"""
        field_descriptions = {}
        for col in df.columns:
            dtype = str(df[col].dtype)
            unique_count = df[col].nunique()
            field_descriptions[col] = f"{col} ({dtype}, {unique_count} unique values)"
        return field_descriptions
    
    async def process_query(self, 
                           question: str,
                           dataset_name: str = "default",
                           chart_type: Optional[str] = None,
                           session_id: Optional[str] = None) -> VisualizationResponse:
        """
        处理查询（使用已加载的数据集）
        
        Args:
            question: 用户问题
            dataset_name: 要使用的数据集名称
            chart_type: 指定图表类型（可选）
            session_id: 会话ID（可选）
            
        Returns:
            可视化结果响应
        """
        if dataset_name not in self.datasets:
            raise ValueError(f"数据集 '{dataset_name}' 未找到。请先使用 load_dataset() 加载数据集。")
        
        dataset = self.datasets[dataset_name]
        
        return await self._process_with_dataset(
            question=question,
            dataset=dataset,
            chart_type=chart_type,
            session_id=session_id
        )
    
    async def visualize(self, 
                       question: str, 
                       data: Union[pd.DataFrame, Dict[str, pd.DataFrame]] = None,
                       dataset_name: str = None,
                       dataset_desc: str = None,
                       table_desc: str = None,
                       field_desc: Dict[str, str] = None,
                       chart_type: Optional[str] = None,
                       session_id: Optional[str] = None) -> VisualizationResponse:
        """
        完整的数据可视化流程
        
        Args:
            question: 用户问题
            data: 数据，可以是单个DataFrame或包含多个DataFrame的字典（可选，如果已加载数据集）
            dataset_name: 数据集名称（如果使用已加载的数据集）
            dataset_desc: 数据集描述（仅在传入data时使用）
            table_desc: 表格描述（仅在传入data时使用）
            field_desc: 字段描述（仅在传入data时使用）
            chart_type: 指定图表类型（可选）
            session_id: 会话ID（可选）
            
        Returns:
            可视化结果响应
        """
        if session_id is None:
            session_id = f"session_{int(time.time() * 1000)}"
        
        # 如果指定了dataset_name，使用已加载的数据集
        if dataset_name and dataset_name in self.datasets:
            return await self.process_query(
                question=question,
                dataset_name=dataset_name,
                chart_type=chart_type,
                session_id=session_id
            )
        
        # 如果传入了data，临时创建数据集
        if data is not None:
            if isinstance(data, dict):
                # 多数据集模式，暂时使用第一个数据集
                temp_dataset_name = list(data.keys())[0]
                df = data[temp_dataset_name]
                self._log_info(f"使用数据集: {temp_dataset_name}", session_id)
            else:
                df = data
                temp_dataset_name = "temp_dataset"
            
            # 创建临时Dataset对象
            dataset = self._create_dataset_from_dataframe(
                df=df, 
                name=temp_dataset_name,
                dataset_desc=dataset_desc,
                table_desc=table_desc,
                field_desc=field_desc
            )
            
            return await self._process_with_dataset(
                question=question,
                dataset=dataset,
                chart_type=chart_type,
                session_id=session_id
            )
        
        # 如果既没有指定dataset_name也没有传入data
        if len(self.datasets) == 1:
            # 如果只有一个数据集，使用它
            dataset_name = list(self.datasets.keys())[0]
            return await self.process_query(
                question=question,
                dataset_name=dataset_name,
                chart_type=chart_type,
                session_id=session_id
            )
        elif len(self.datasets) > 1:
            raise ValueError("存在多个数据集，请指定 dataset_name 参数")
        else:
            raise ValueError("没有可用的数据集，请先使用 load_dataset() 加载数据或传入 data 参数")
    
    async def _process_with_dataset(self,
                                   question: str,
                                   dataset: Dataset,
                                   chart_type: Optional[str] = None,
                                   session_id: Optional[str] = None) -> VisualizationResponse:
        """使用指定数据集处理查询"""
        start_time = time.time()
        
        try:
            self._log_info("开始可视化流程", session_id)
            
            # 优化判断：检查是否可以使用快速处理
            data_rows = len(dataset.df)
            self._log_info(f"数据行数: {data_rows}", session_id)
            
            # 场景1：单行数据直接返回table
            if data_rows == 1:
                self._log_info("检测到单行数据，使用快速table处理", session_id)
                return await self._process_single_row_data(dataset, question, session_id, start_time)
            
            # 场景2：小数据集且不需要分组过滤时直接生成图表
            if (data_rows < self.small_data_threshold and 
                not self.enable_grouping_filtering):
                self._log_info(f"检测到小数据集({data_rows}行)且无需分组过滤，使用快速图表生成", session_id)
                return await self._process_small_dataset_direct(dataset, question, chart_type, session_id, start_time)
            
            # 标准流程处理
            self._log_info("使用标准完整流程处理", session_id)
            
            # 创建查询上下文
            context = QueryContext(
                question=question,
                language="zh",
                user_preferences={}
            )
            
            # 阶段1: 图表需求分析
            self._log_info("阶段1: 开始图表需求分析", session_id)
            analysis_result = await self.chart_identifier.analyze_chart_requirement(
                dataset, context
            )
            
            # 如果指定了图表类型，覆盖分析结果
            if chart_type:
                analysis_result.chart_type = chart_type
                analysis_result.need_chart = chart_type != "table"
            
            # 阶段2: 数据处理
            self._log_info("阶段2: 开始数据处理", session_id)
            processing_result = await self.data_processor.process_data(
                dataset, analysis_result
            )
            
            # 阶段3: 输出生成
            self._log_info("阶段3: 开始输出生成", session_id)
            visualization_result = await self.output_generator.generate_visualization(
                dataset, analysis_result, processing_result
            )
            
            execution_time = time.time() - start_time
            self._log_info(f"操作 '完整可视化流程' 执行时间: {execution_time:.3f}秒", session_id)
            
            return VisualizationResponse(
                success=visualization_result.success,
                chart_config=visualization_result.config,
                data=processing_result.processed_df.to_dict('records') if not processing_result.processed_df.empty else [],
                chart_type=analysis_result.chart_type,
                execution_time=execution_time,
                session_id=session_id,
                analysis_result=analysis_result.to_dict(include_debug_info=self.debug_mode),
                message="可视化完成" if visualization_result.success else visualization_result.error_message,
                error=visualization_result.error_message if not visualization_result.success else None
            )
            
        except Exception as e:
            execution_time = time.time() - start_time
            error_msg = f"可视化流程失败: {str(e)}"
            self._log_error(error_msg, session_id)
            
            return VisualizationResponse(
                success=False,
                chart_config=None,
                data=[],
                chart_type=chart_type,
                execution_time=execution_time,
                session_id=session_id,
                error=error_msg,
                message=error_msg
            )
    
    def _create_dataset_from_dataframe(self, 
                                      df: pd.DataFrame, 
                                      name: str = "default",
                                      dataset_desc: str = None,
                                      table_desc: str = None,
                                      field_desc: Dict[str, str] = None) -> Dataset:
        """从DataFrame创建Dataset对象"""
        # 使用提供的描述或生成默认描述
        if dataset_desc is None:
            dataset_desc = f"数据集 {name}"
        
        if table_desc is None:
            table_desc = f"包含 {len(df)} 行 {len(df.columns)} 列的数据表"
        
        if field_desc is None:
            field_desc = self._generate_default_field_descriptions(df)
        else:
            # 补充缺失的字段描述
            complete_field_desc = self._generate_default_field_descriptions(df)
            complete_field_desc.update(field_desc)
            field_desc = complete_field_desc
        
        return Dataset(
            dataset_id=f"dataset_{name}_{int(time.time())}",
            df=df,
            dataset_description=dataset_desc,
            table_description=table_desc,
            field_descriptions=field_desc
        )
    
    def clear_datasets(self) -> None:
        """清除所有已加载的数据集"""
        self.datasets.clear()
        self._log_info("所有数据集已清除")
    
    def list_datasets(self) -> List[str]:
        """列出所有已加载的数据集名称"""
        return list(self.datasets.keys())
    
    def get_dataset_info(self, dataset_name: str) -> Dict[str, Any]:
        """获取数据集信息"""
        if dataset_name not in self.datasets:
            raise ValueError(f"数据集 '{dataset_name}' 未找到")
        
        dataset = self.datasets[dataset_name]
        return {
            "dataset_name": dataset_name,
            "dataset_description": dataset.dataset_description,
            "table_description": dataset.table_description,
            "field_descriptions": dataset.field_descriptions,
            "shape": dataset.df.shape,
            "columns": list(dataset.df.columns)
        }
    
    async def analyze_only(self, 
                          question: str, 
                          data: Union[pd.DataFrame, Dict[str, pd.DataFrame]],
                          chart_type: Optional[str] = None,
                          session_id: Optional[str] = None) -> Dict[str, Any]:
        """
        仅进行需求分析，不处理数据
        
        Args:
            question: 用户问题
            data: 数据
            chart_type: 指定图表类型（可选）
            session_id: 会话ID（可选）
            
        Returns:
            分析结果
        """
        if session_id is None:
            session_id = f"session_{int(time.time() * 1000)}"
        
        try:
            # 处理数据输入
            if isinstance(data, dict):
                dataset_name = list(data.keys())[0]
                df = data[dataset_name]
            else:
                df = data
                dataset_name = "default"
            
            # 创建Dataset对象
            dataset = self._create_dataset_from_dataframe(df, dataset_name)
            
            # 创建查询上下文
            context = QueryContext(
                question=question,
                language="zh",
                user_preferences={}
            )
            
            analysis_result = await self.chart_identifier.analyze_chart_requirement(
                dataset, context
            )
            
            return {
                "success": True,
                "analysis": analysis_result.__dict__,
                "session_id": session_id
            }
            
        except Exception as e:
            self._log_error(f"需求分析失败: {str(e)}", session_id)
            return {
                "success": False,
                "error": str(e),
                "session_id": session_id
            }
    
    def get_supported_chart_types(self) -> List[str]:
        """
        获取支持的图表类型
        
        Returns:
            支持的图表类型列表
        """
        return self.config_manager.get_all_chart_types()
    
    def get_engine_info(self) -> Dict[str, Any]:
        """
        获取引擎信息
        
        Returns:
            引擎信息字典
        """
        return {
            "name": "DataEngine",
            "version": "1.0.0",
            "supported_chart_types": self.get_supported_chart_types(),
            "llm_model": self.llm_client.model,
            "description": "基于LLM的智能数据可视化引擎"
        }